#!/usr/bin/env python

# *************************************************************************

# Copyright 2010 Texas A&M University, 
#                College Station, Texas
#                Contact: R Crosby (rwc@cs.tamu.edu)
#
# Copyright 2003, 2004 by BiRC -- Bioinformatics Research Center
#                University of Aarhus, Denmark
#                Contact: Thomas Mailund <mailund@birc.dk>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

# *************************************************************************

import os
import re

# *************************************************************************

def SetupWorkDirectory ( config ):

    try:
        workdir = config.get ( 'directories', 'work' )
    except ConfigParmer.NoOptionError, err:
        print str(err)
        sys.exit(2)

    # get a name to use in the workdir
    pid = os.getpid()
    piddir = workdir + '/' + str(pid)

    # make the pid directory, create if required
    os.makedirs(piddir)

    return piddir

# *************************************************************************

def GetSegmentFile(piddir, treeno, opt):

    fn = 'trees_' + '%09d' % treeno + '.newick'
    tf = piddir + '/' + fn
    
    sf = open(tf, 'w')

    if opt['verbose']:
        print 'Tree file:', tf

    return sf, fn

# *************************************************************************

def PartitionTrees ( piddir, tfiles, opt ):

    ntrees = 0                                    # Total trees output 
    limhit = False                                # Loop termination control
    nsegments = 0                                 # Number of segments generated
    segfiles = []                                 # Names of the segment files
    

    for fl in tfiles:                             # Process through the tree files specified

        with open(fl) as fp:
        
            buf = ''

            for line in fp:                       # Process all the lines in the buffer

                line = line.rstrip()              # Trim newlined
                trees = re.split('([^;]*;)', line)# Split into trees

                for seg in trees:                 # Process segments found
                    if seg == '':                 # Ignore empty segments
                        continue

                    if re.search(';$', seg):      # Does a tree end here?

                        if not ntrees % opt['segsize']:  # Need a new segment file?
                            nsegments += 1
                            segfile,  fn = GetSegmentFile(piddir, ntrees, opt)
                            segfiles.append(fn)

                        buf = buf + seg
                        ntrees += 1
                        segfile.write(buf+'\n')
                        buf = ''

                    else:                         # No end of tree, save...

                        buf = buf + seg           # Add to accumulated tree

                    if opt['tlimit'] and ntrees >= opt['tlimit']: # At max number of trees
                        limhit = True
                        break

                if limhit:
                    break

        if limhit:
            break

    segfile.close

    return ntrees, segfiles

# *************************************************************************

if __name__ == "__main__":
    print 'Needs to be run from QQFront main'
